{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 08: modify a linear system on restriction\n",
    "\n",
    "In this tutorial we solve the problem\n",
    "\n",
    "$$\\begin{cases}\n",
    "-\\Delta u = f, & \\text{in } \\Omega,\\\\\n",
    " u   = g, & \\text{on } \\Gamma = \\partial\\Omega,\n",
    "\\end{cases}$$\n",
    "\n",
    "where $\\Omega$ is the unit ball in 2D.\n",
    "\n",
    "We compare the following two cases:\n",
    "* **strong imposition of Dirichlet BCs**:\n",
    "the corresponding weak formulation is\n",
    "$$\n",
    "\\text{find } u \\in V_g \\text{ s.t. } \\int_\\Omega \\nabla u \\cdot \\nabla v = \\int_\\Omega f v, \\quad \\forall v \\in V_0\\\\\n",
    "$$\n",
    "where\n",
    "$$\n",
    "V_g = \\{v \\in H^1(\\Omega): v|_\\Gamma = g\\},\\\\\n",
    "V_0 = \\{v \\in H^1(\\Omega): v|_\\Gamma = 0\\}.\\\\\n",
    "$$\n",
    "* **penalty imposition of Dirichlet BCs**: this requires first the discretization of the associated homogeneous Neumann problem\n",
    "$$\n",
    "\\text{find } u \\in V \\text{ s.t. } \\int_\\Omega \\nabla u \\cdot \\nabla v = \\int_\\Omega f v, \\quad \\forall v \\in V\\\\\n",
    "$$\n",
    "where\n",
    "$$\n",
    "V = H^1(\\Omega),\n",
    "$$\n",
    "which we rewrite in matrix form as\n",
    "$$ \\text{find } x \\in \\mathbb{R}^n \\text{ s.t. } A x = b,$$\n",
    "where $A \\in \\mathbb{R}^{n \\times n}$ and $b \\in \\mathbb{R}^n$ are the left-hand side and right-hand side of the discrete Neumann problem.\n",
    "\n",
    "In order to impose Dirichlet boundary conditions, we solve the modified system\n",
    "$$ \\text{find } \\tilde{x} \\in \\mathbb{R}^n \\text{ s.t. } \\tilde{A} \\tilde{x} = \\tilde{b},$$\n",
    "where\n",
    "$$\\tilde{A} = A + P, \\qquad \\tilde{b} = b + q.$$\n",
    "\n",
    "The penalty matrix $P$ is defined as\n",
    "$$\n",
    "P_{ij} =\n",
    "\\begin{cases}\n",
    "\\mu, & \\text{if $i = j$, and $i$ is a boundary DOF},\\\\\n",
    "0, & \\text{otherwise},\n",
    "\\end{cases}$$\n",
    "and the penalty vector $q$ as\n",
    "$$\n",
    "q_i = \n",
    "\\begin{cases}\n",
    "\\mu g(\\mathbf{x}_i), & \\text{if $i$ is a boundary DOF},\\\\\n",
    "0, & \\text{otherwise},\n",
    "\\end{cases}$$\n",
    "where $\\mu \\in \\mathbb{R}$ is (large) scalar number, and $\\mathbf{x}_i$ denotes the coordinate of the entity associated at DOF $i$.\n",
    "\n",
    "The preferred way to impose non-homogeneous Dirichlet boundary conditions should still\n",
    "either be through DirichletBC objects (see tutorial 01) or lagrange multipliers\n",
    "(see tutorial 03). This example is meant to show how to use the list of degrees of freedom\n",
    "associated to a specific restriction to perform local modifications to assembled\n",
    "matrices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from mpi4py import MPI\n",
    "from petsc4py import PETSc\n",
    "from ufl import grad, inner, Measure, TestFunction, TrialFunction\n",
    "from dolfinx import DirichletBC, Function, FunctionSpace\n",
    "from dolfinx.cpp.la import create_vector\n",
    "from dolfinx.fem import LinearProblem, locate_dofs_topological\n",
    "from dolfinx.io import XDMFFile\n",
    "from dolfinx.plot import create_vtk_topology\n",
    "from multiphenicsx.fem import assemble_matrix, assemble_scalar, assemble_vector, create_matrix\n",
    "import pyvista"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with XDMFFile(MPI.COMM_WORLD, \"data/circle.xdmf\", \"r\") as infile:\n",
    "    mesh = infile.read_mesh()\n",
    "    mesh.topology.create_connectivity_all()\n",
    "    subdomains = infile.read_meshtags(mesh, name=\"subdomains\")\n",
    "    boundaries = infile.read_meshtags(mesh, name=\"boundaries\")\n",
    "facets_Gamma = boundaries.indices[boundaries.values == 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define associated measures\n",
    "dx = Measure(\"dx\")(subdomain_data=subdomains)\n",
    "ds = Measure(\"ds\")(subdomain_data=boundaries)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dolfinx_to_pyvista_mesh(mesh):\n",
    "    num_cells = mesh.topology.index_map(mesh.topology.dim).size_local\n",
    "    cell_entities = np.arange(num_cells, dtype=np.int32)\n",
    "    pyvista_cells, cell_types = create_vtk_topology(mesh, mesh.topology.dim, cell_entities)\n",
    "    grid = pyvista.UnstructuredGrid(pyvista_cells, cell_types, mesh.geometry.x)\n",
    "    return grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pyvista_mesh_plot(mesh):\n",
    "    grid = dolfinx_to_pyvista_mesh(mesh)\n",
    "    plotter = pyvista.PlotterITK()\n",
    "    plotter.add_mesh(grid)\n",
    "    plotter.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_mesh_plot(mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_penalty_system(V, restriction, penalty, g):\n",
    "    \"\"\"\n",
    "    Generate matrix and vector to be added to the system to handle the penalty terms\n",
    "    \"\"\"\n",
    "    # Assemble penalty matrix.\n",
    "    p = u * v * ds  # this form is not really used per se, just to allocate the correct sparsity pattern\n",
    "    P = create_matrix(p)  # this creates a zero matrix\n",
    "    for dof in restriction:\n",
    "        P.setValuesLocal([dof], [dof], [penalty])  # set P_{ij} = penalty * \\delta_{ij}, if i on boundary\n",
    "    P.assemble()\n",
    "    # Next, assemble the penalty vector.\n",
    "    q = create_vector(V.dofmap.index_map, V.dofmap.index_map_bs)\n",
    "    with q.localForm() as q_local, g.vector.localForm() as g_local:\n",
    "        q_local[restriction] = penalty * g_local[restriction]\n",
    "    # Return matrix and vector\n",
    "    return P, q"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Penalty imposition of Dirichlet BCs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function space\n",
    "V = FunctionSpace(mesh, (\"Lagrange\", 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define trial and test functions\n",
    "u = TrialFunction(V)\n",
    "v = TestFunction(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define problem forms\n",
    "g = Function(V)\n",
    "g.interpolate(lambda x: np.sin(3 * x[0] + 1) * np.sin(3 * x[1] + 1))\n",
    "a = inner(grad(u), grad(v)) * dx\n",
    "f = v * dx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assemble the linear system\n",
    "A = assemble_matrix(a)\n",
    "A.assemble()\n",
    "F = assemble_vector(f)\n",
    "F.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add penalty terms\n",
    "dofs_V_Gamma = locate_dofs_topological(V, boundaries.dim, facets_Gamma)\n",
    "(P, Q) = generate_penalty_system(V, dofs_V_Gamma, 1.e10, g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u = Function(V)\n",
    "ksp = PETSc.KSP()\n",
    "ksp.create(mesh.mpi_comm())\n",
    "ksp.setOperators(A + P)\n",
    "ksp.setType(\"preonly\")\n",
    "ksp.getPC().setType(\"lu\")\n",
    "ksp.getPC().setFactorSolverType(\"mumps\")\n",
    "ksp.solve(F + Q, u.vector)\n",
    "u.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pyvista_scalar_field_plot(mesh, scalar_field, name):\n",
    "    grid = dolfinx_to_pyvista_mesh(mesh)\n",
    "    grid.point_arrays[name] = scalar_field.compute_point_values()\n",
    "    grid.set_active_scalars(name)\n",
    "    plotter = pyvista.PlotterITK()\n",
    "    plotter.add_mesh(grid)\n",
    "    plotter.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, u, \"u\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Strong imposition of Dirichlet BCs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define Dirichlet BC object on Gamma\n",
    "bc_ex = DirichletBC(g, dofs_V_Gamma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "u_ex = Function(V)\n",
    "problem_ex = LinearProblem(\n",
    "    a, f, bcs=[bc_ex], u=u_ex,\n",
    "    petsc_options={\"ksp_type\": \"preonly\", \"pc_type\": \"lu\", \"pc_factor_mat_solver_type\": \"mumps\"})\n",
    "problem_ex.solve()\n",
    "u_ex.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, u_ex, \"u\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparison and error computation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_ex_norm = np.sqrt(mesh.mpi_comm().allreduce(assemble_scalar(inner(grad(u_ex), grad(u_ex)) * dx), op=MPI.SUM))\n",
    "err_norm = np.sqrt(mesh.mpi_comm().allreduce(assemble_scalar(inner(grad(u_ex - u), grad(u_ex - u)) * dx), op=MPI.SUM))\n",
    "print(\"Relative error is equal to\", err_norm / u_ex_norm)\n",
    "assert np.isclose(err_norm / u_ex_norm, 0., atol=1.e-10)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython"
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
